from data.mnist_dataset import MNISTDataModule
from torch.utils.data import Subset
import numpy as np

class MNISTTestDataModule(MNISTDataModule):
    """测试专用数据集：支持按比例采样，继承自原有数据集类"""
    def __init__(self, batch_size=32, data_dir="./data", sample_ratio=1.0):
        super().__init__(batch_size=batch_size, data_dir=data_dir)
        self.sample_ratio = sample_ratio  # 新增采样比例参数

    def setup(self, stage=None):
        # 调用父类方法加载完整数据
        super().setup(stage)
        # 仅在测试模式（sample_ratio < 1）时采样
        if self.sample_ratio < 1.0:
            # 训练集采样
            train_size = int(len(self.train_dataset) * self.sample_ratio)
            train_indices = np.random.choice(len(self.train_dataset), train_size, replace=False)
            self.train_dataset = Subset(self.train_dataset, train_indices)
            # 验证集采样
            val_size = int(len(self.val_dataset) * self.sample_ratio)
            val_indices = np.random.choice(len(self.val_dataset), val_size, replace=False)
            self.val_dataset = Subset(self.val_dataset, val_indices)
            print(f"⚠️ 测试模式：使用{self.sample_ratio*100}%数据（训练集{len(self.train_dataset)}，验证集{len(self.val_dataset)}）")